#!/usr/bin/env python3
"""Plot and summarize MPGP terminal error statistics."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Tuple

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402


def load_mpgp_result(npz_path: Path) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Load goal error, min-distance, and cost arrays from an MPGP/MPC result file."""
    if not npz_path.exists():
        raise FileNotFoundError(f"NPZ file not found: {npz_path}")
    data = np.load(npz_path, allow_pickle=False)
    if "error_norms" not in data or "min_distances" not in data:
        raise KeyError("NPZ missing required arrays: error_norms and min_distances")
    goal_errors = np.asarray(data["error_norms"], dtype=np.float64)
    min_distances = np.asarray(data["min_distances"], dtype=np.float64)
    costs = np.asarray(data["costs"], dtype=np.float64) if "costs" in data else None
    return goal_errors, min_distances, costs


def compute_stats(goal_errors: np.ndarray, min_distances: np.ndarray, costs: np.ndarray | None) -> dict:
    """Compute per-step and aggregate statistics."""
    per_step_mean = np.nanmean(goal_errors, axis=0)
    per_step_var = np.nanvar(goal_errors, axis=0)
    overall_goal_error_mean = float(np.nanmean(goal_errors))
    overall_goal_error_var = float(np.nanvar(goal_errors))
    min_distance_mean = float(np.nanmean(min_distances))
    min_distance_var = float(np.nanvar(min_distances))
    # 成本统计：先按 episode 求总成本，再对总成本求均值/方差；忽略 NaN/0，并滤掉高异常值
    if costs is not None:
        valid = np.isfinite(costs) & (costs != 0)
        costs_clean = np.where(valid, costs, np.nan)
        episode_costs = np.nansum(costs_clean, axis=1)  # [n_episodes]
        finite_costs = episode_costs[np.isfinite(episode_costs)]
        if finite_costs.size == 0:
            cost_mean = None
            cost_var = None
        else:
            cutoff = np.nanpercentile(finite_costs, 95)
            filtered = finite_costs[finite_costs <= cutoff]
            # 计算方差时剔除大于 3000 的 episode 成本
            filtered = filtered[filtered <= 3000]
            if filtered.size == 0:
                cost_mean = None
                cost_var = None
            else:
                cost_mean = float(np.nanmean(filtered))
                cost_var = float(np.nanvar(filtered))
    else:
        cost_mean = None
        cost_var = None
    return {
        "per_step_mean": per_step_mean,
        "per_step_var": per_step_var,
        "overall_goal_error_mean": overall_goal_error_mean,
        "overall_goal_error_var": overall_goal_error_var,
        "min_distance_mean": min_distance_mean,
        "min_distance_var": min_distance_var,
        "cost_mean": cost_mean,
        "cost_var": cost_var,
    }


def plot_per_step_mean(
    per_step_mean: np.ndarray,
    per_step_var: np.ndarray,
    output_path: Path,
) -> None:
    """Plot mean terminal error per MPC step with a ±1 std band."""
    steps = np.arange(1, per_step_mean.shape[0] + 1)
    std = np.sqrt(np.clip(per_step_var, a_min=0.0, a_max=None))
    valid = np.isfinite(per_step_mean) & np.isfinite(std)

    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.plot(steps, per_step_mean, marker="o", color="tab:blue", linewidth=2, label="Mean terminal error")
    if valid.any():
        ax.fill_between(
            steps[valid],
            per_step_mean[valid] - std[valid],
            per_step_mean[valid] + std[valid],
            color="tab:blue",
            alpha=0.2,
            label="±1 std",
        )
    ax.set_xlabel("MPC step")
    ax.set_ylabel("Terminal error")
    ax.set_title("Per-step mean terminal error (MPGP)")
    ax.grid(True, linestyle="--", alpha=0.6)
    ax.legend()
    fig.tight_layout()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, dpi=300)
    plt.close(fig)


def save_summary(summary: dict, output_path: Path, source_path: Path, goal_errors: np.ndarray) -> None:
    """Persist key statistics (overall goal error mean, min-distance stats, cost stats)."""
    output_path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "source_file": str(source_path),
        "n_episodes": int(goal_errors.shape[0]),
        "n_mpc_steps": int(goal_errors.shape[1]),
        "overall_goal_error_mean": summary["overall_goal_error_mean"],
        "overall_goal_error_var": summary["overall_goal_error_var"],
        "min_distance_mean": summary["min_distance_mean"],
        "min_distance_var": summary["min_distance_var"],
    }
    if summary.get("cost_mean") is not None:
        payload["cost_mean"] = summary["cost_mean"]
        payload["cost_var"] = summary["cost_var"]
    with output_path.open("w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Plot per-step MPGP terminal error mean and summarize metrics.",
    )
    parser.add_argument(
        "--input",
        type=Path,
        default=Path(r"D:\test8\multi-uav-trajectory-planning\results\demo_mpc\mpc_result.npz"),
        help="Path to the NPZ file containing error_norms and min_distances.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("plot"),
        help="Directory to write plots and summary.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    goal_errors, min_distances, costs = load_mpgp_result(args.input)
    stats = compute_stats(goal_errors, min_distances, costs)

    plot_path = args.output_dir / "mpc_terminal_error_mean.png"
    plot_per_step_mean(stats["per_step_mean"], stats["per_step_var"], plot_path)

    summary_path = args.output_dir / "mpc_metrics_summary.json"
    save_summary(stats, summary_path, args.input, goal_errors)

    print(f"Saved per-step mean plot to {plot_path}")
    print(f"Saved summary metrics to {summary_path}")


if __name__ == "__main__":
    main()
